"""
用 cqlib 从（目前：非负实）向量制备量子态 |x> 的电路生成器（带中文注释）

总体思路（“二叉树幅度装载”）：
--------------------------------
将长度为 L 的输入向量 x（目前限定为非负实数）补齐到 2^n 维，并进行归一化。
把归一化后的幅值按二叉树自顶向下地划分为“左子/右子”的能量（幅值平方和）比例，
在量子电路上通过一系列（多控）RY 旋转，将“分支比例”刻到相应的位上：
- 根节点：在最低有效位（LSB）上做一次 RY(theta0)，把概率在左右子树之间按比例分配；
- 非叶子节点：对更深一层的比特位，受控（多控）RY 进一步细分概率到更小分支；
- 叶子节点：如果需要处理复相位对（left/right 的相位不同），在相应目标位上额外加一个 RZ（相对相位）。

注：本文件当前实现中 `prepare_state_from_vector` 对输入做了约束：
    仅接受“元素非负的实向量”。这样叶子相位只会是 (0, π) 的若干离散情况，
    复杂相位对（非 0/π）会通过 RZ 做相对相位补偿。若你要真正支持“任意复向量”，
    需要移除该约束，并在全树层面系统性加载相位（不仅是叶子对），这通常涉及
    额外的受控 RZ/CP 架构或者分层的相位同步策略。

多控 RY 的实现：
----------------
k>1 时采用线性 Toffoli 链把所有控制位“与”到辅助位（ancillas）上，
再用一个 CRY 完成受控旋转，然后逆向 uncompute 恢复辅助位到 |0>。
所需辅助位数量约为 (k-1)。

模拟验证：
----------
附带一个简易的仿真助手 `simulate_statevector`（优先 SimpleSimulator，
备选 StatevectorSimulator），并提供 run_case 测试样例：
对比“电路输出的态矢量”与“目标态（含 ancilla 在 |0...0>）”的 Fidelity。

注意：
-----
- 为了与“二叉树幅度”一致，我们对量子比特采取 LSB 为最底层分裂位的布置，并在
  开始时用 swap 将“数据位”从 MSB/LSB 顺序调整到预期的物理布局。
- 如果你已有既定的比特位序规范，可以把 swap 段删去或改写。

"""

from __future__ import annotations
from dataclasses import dataclass
from typing import Optional, List, Tuple

import numpy as np

# 为避免导入重量级依赖，这里按你的路径使用 Algorithm_Circuit 作为 Circuit
# 若你在实际工程中已有 from cqlib.circuits import Circuit，可替换为官方导入
from circuits.algorithm_circuit import Algorithm_Circuit as Circuit


# =============================================================================
# 一、二叉树结构：用于自顶向下记录“子树的幅值平方和（及叶子相位）”
# =============================================================================
class BinaryTree:
    """
    二叉树的每个节点保存一个“值”：
    - 对非叶子：value 是“该子树幅值平方和 s ∈ [0,1]”
    - 对叶子：value 是长度为 2 的列表 [s, phase]，其中 s=|a|^2，phase=arg(a)，a 为叶子对应幅度
    同时记录：
    - level：当前节点在树中的层数（根为 0）
    - branch：从根到该节点的分支编号（把左右分支视为二进制 0/1 累计）
    """
    def __init__(self, value, level=0, branch=0):
        self.value = value
        self.child_left: Optional['BinaryTree'] = None
        self.child_right: Optional['BinaryTree'] = None
        self.level = level
        self.branch = branch

    def insert_left(self, value, branch):
        """插入/前插左孩子；若已有左孩子，则做“头插法”并维护层号与分支号。"""
        if self.child_left is None:
            self.child_left = BinaryTree(value, self.level + 1, branch)
        else:
            new_node = BinaryTree(value, self.level + 1, branch)
            # 原左孩子整体下移一层；分支号×2（相当于在二进制末尾补 0）
            self.child_left.level += 1
            self.child_left.branch *= 2
            new_node.child_left = self.child_left
            self.child_left = new_node

    def insert_right(self, value, branch):
        """插入/前插右孩子；若已有右孩子，则做“头插法”并维护层号与分支号。"""
        if self.child_right is None:
            self.child_right = BinaryTree(value, self.level + 1, branch)
        else:
            new_node = BinaryTree(value, self.level + 1, branch)
            # 原右孩子整体下移一层；分支号×2+1（相当于在二进制末尾补 1）
            self.child_right.level += 1
            self.child_right.branch = self.child_right.branch * 2 + 1
            new_node.child_right = self.child_right
            self.child_right = new_node

    def get_value(self):
        return self.value
    def get_left(self): return self.child_left
    def get_right(self): return self.child_right
    def get_level(self): return self.level
    def get_branch(self): return self.branch


# =============================================================================
# 二、构建“幅度+相位”的二叉树
# =============================================================================
def create_tree_from_vector(x: List[complex]) -> Tuple[float, BinaryTree]:
    """
    将输入向量 x 归一化并构造二叉树：
    - 若 len(x) 不是 2^n，先在末尾用 0 补齐；
    - 根节点为 1.0（表示总概率和 1），叶子节点保存 [s, phase]；
    - 非叶子节点保存子树概率和 s_left / s_right，递归构造。

    返回：
        (原始范数 ||x||, 二叉树根节点)
    """
    x = np.asarray(x, dtype=complex)
    # 1) 补齐到 2^n 维
    n = int(np.ceil(np.log2(len(x)))) if len(x) > 1 else 1
    N = 1 << n
    if len(x) < N:
        x = np.concatenate([x, np.zeros(N - len(x), dtype=complex)])

    # 2) 归一化
    norm = np.linalg.norm(x)
    if norm == 0:
        raise ValueError("Input vector must be non-zero.")
    u_x = x / norm

    # 3) 特例：单元素（n=0 的退化情况）
    if len(u_x) == 1:
        tree = BinaryTree([1.0, float(np.angle(u_x[0]))])
        return float(norm), tree

    # 4) 递归构造二叉树：根为 1.0
    tree = BinaryTree(1.0, level=0, branch=0)
    mid = len(u_x) // 2
    create_tree_recursively(u_x[:mid], tree, "left", tree.get_branch())
    create_tree_recursively(u_x[mid:], tree, "right", tree.get_branch())
    return float(norm), tree


def create_tree_recursively(x: np.ndarray, tree: BinaryTree, side: str, branch: int):
    """
    对数组 x（已是 2^k 长度）：
    - 如果 len(x)==1，为叶子：记录 [s, phase]；
    - 否则为中间节点：记录 s = sum |x|^2，并继续对子数组左右递归。
    - 分支号 branch 在递归时按“左 0、右 1”的规则逐步更新。
    """
    s = float(np.sum(np.abs(x) ** 2))
    if len(x) == 1:
        phase = float(np.angle(x[0]))
        s = [s, phase]

    if side == "left":
        tree.insert_left(s, branch * 2)
    elif side == "right":
        tree.insert_right(s, branch * 2 + 1)
    else:
        raise ValueError("side must be 'left' or 'right'")

    if len(x) > 1:
        mid = len(x) // 2
        x_left = x[:mid]
        x_right = x[mid:]

        # 边界补齐：当一侧已是叶子而另一侧仍可分裂时，补 0 防止不平衡
        if len(x_left) == 1 and len(x_right) > 1:
            x_left = np.array([x_left[0], 0], dtype=complex)
        elif len(x_right) == 1 and len(x_left) > 1:
            x_right = np.array([x_right[0], 0], dtype=complex)

        child = tree.get_left() if side == "left" else tree.get_right()
        create_tree_recursively(x_left, child, "left", child.get_branch())
        create_tree_recursively(x_right, child, "right", child.get_branch())


# =============================================================================
# 三、若干电路构造辅助函数
# =============================================================================
def get_binary(number: int, integer_precision: int = 1) -> str:
    """把整数 number 转为定长（integer_precision 位）二进制字符串。"""
    return format(int(number), '0{}b'.format(integer_precision))


def _apply_control_pattern(circ: Circuit, controls: List[int], pattern_bits: str):
    """
    为了让多控门对“特定的比特模式”触发，这里先把控制位中需要匹配 0 的位置用 X 翻转，
    从而把“匹配 pattern_bits”的配置统一转换为“全 1”。
    返回：本次翻转过的控制位索引列表，用于之后恢复。
    """
    flipped = []
    for idx, bit in zip(controls, pattern_bits):
        if bit == '0':
            circ.x(idx)
            flipped.append(idx)
    return flipped


def _undo_control_pattern(circ: Circuit, flipped: List[int]):
    """将前面为匹配 0 位而做的 X 翻转逐一撤销。"""
    for idx in flipped:
        circ.x(idx)


def _mc_ry(circ: Circuit, controls: List[int], target: int, theta: float, ancillas: List[int]):
    """
    实现“多控制位”的 RY(theta)：
    - k=0：直接在 target 上 RY
    - k=1：直接 CRY
    - k>1：用 Toffoli 链把所有控制位的“与”写到 ancillas 上，再用一个 CRY 驱动 target，
            之后逆向把 ancillas uncompute 回 |0...0>

    参数：
        controls: 控制位索引（从高到低或从低到高由调用者统一约定）
        target:   目标位索引
        ancillas: 可用辅助位索引列表（长度至少 k-1）
    """
    k = len(controls)
    if k == 0:
        circ.ry(target, theta)
        return
    if k == 1:
        circ.cry(controls[0], target, theta)
        return
    if len(ancillas) < k - 1:
        raise ValueError(f"Need at least {k-1} ancillas for {k}-controlled RY; got {len(ancillas)}")

    # 1) 正向：把 controls 依次 AND 到 ancillas 链上
    used = []
    # 先把前两个控制位 AND 到第一个辅助位
    circ.ccx(controls[0], controls[1], ancillas[0]); used.append((controls[0], controls[1], ancillas[0]))
    # 再把后续控制位与前一个“累积位”继续 AND
    for i in range(2, k):
        circ.ccx(ancillas[i-2], controls[i], ancillas[i-1])
        used.append((ancillas[i-2], controls[i], ancillas[i-1]))

    # 2) 单控 RY：由“最终累积位”控制
    final_ctrl = ancillas[k-2]
    circ.cry(final_ctrl, target, theta)

    # 3) 逆向：撤销 Toffoli 链，恢复 ancillas=|0...0>
    for c1, c2, tgt in reversed(used):
        circ.ccx(c1, c2, tgt)


def _process_node(circ: Circuit, first_qubit: int, n: int, u_x: BinaryTree, ancillas: List[int]):
    """
    递归地处理某个子树 u_x：
    - 计算“当前节点左子树相对于父节点的概率占比”，得到旋转角 theta；
    - 根据分支号 branch，构造控制位模式并在相应目标位上施加（多控）RY；
    - 若叶子两侧相位不是 {0,π} 离散组合，则在目标位上补一个相对相位的 RZ；
    - 继续深入处理左右子树。
    """
    u_l = u_x.get_left()
    u_r = u_x.get_right()
    if u_x.get_value() == 0 or u_l is None or u_r is None:
        return

    # 计算旋转角：theta = 2 * arccos( sqrt( s_left / s_parent ) )
    u_l_is_leaf = (u_l.get_left() is None)
    node_value = u_l.get_value()[0] if u_l_is_leaf else u_l.get_value()
    theta = 2 * np.arccos(np.sqrt(node_value / u_x.get_value()))

    # 叶子相位处理：若是非 {0, π} 的复相位对，则记号以便稍后加 RZ
    is_complex_pair = False
    if u_l_is_leaf:
        phase_l, phase_r = u_l.get_value()[1], u_r.get_value()[1]
        if (phase_l, phase_r) not in [(0.0, 0.0), (0.0, np.pi), (np.pi, 0.0), (np.pi, np.pi)]:
            is_complex_pair = True
        # 针对 {0,π} 的组合，可把符号吸收到 theta 的方向上，减少 RZ
        elif (phase_l, phase_r) == (0.0, np.pi):
            theta = -theta
        elif (phase_l, phase_r) == (np.pi, 0.0):
            theta = np.pi - theta
        elif (phase_l, phase_r) == (np.pi, np.pi):
            theta = np.pi + theta

    # 控制位是“之前已经决定好的更高层比特”
    depth = u_l.get_level() - 1
    # 控制位按 LSB→MSB 的物理顺序取出（first_qubit + (n-1-d)）
    controls = [first_qubit + (n - 1 - d) for d in range(depth)]
    pattern_bits = get_binary(u_x.get_branch(), depth)
    flipped = _apply_control_pattern(circ, controls, pattern_bits)

    # 目标位是“当前这层的物理位”（同样按 LSB 优先的布局）
    target = first_qubit + (n - 1 - (u_l.get_level() - 1))
    _mc_ry(circ, controls, target, float(theta), ancillas)

    # 撤销为匹配 0 位而做的 X 翻转
    _undo_control_pattern(circ, flipped)

    # 若叶子是“复杂相位对”，补一个相对相位的 RZ（仅做左右两叶的相对）
    if u_l_is_leaf and is_complex_pair:
        circ.rz(target, float(u_r.get_value()[1] - u_l.get_value()[1]))

    # 继续递归左右子树（若左孩子不是叶子）
    if not u_l_is_leaf:
        # 可选：放一个 barrier 便于可视化分段
        used = [first_qubit + (n - 1 - d) for d in range(u_l.get_level())]
        circ.barrier(*used)
        _process_node(circ, first_qubit, n, u_l, ancillas)
        _process_node(circ, first_qubit, n, u_r, ancillas)


# =============================================================================
# 四、主入口：从（目前：非负实）向量构造状态制备电路
# =============================================================================
def prepare_state_from_vector(x: List[complex],
                              pos_init: int = 0) -> Tuple[Circuit, float, BinaryTree, List[int]]:
    """
    构造一个在 n（数据）比特 + (n-2)（辅助）比特上制备 |x>/||x|| 的电路。
    注意：当前实现限制 x 为“元素非负的实向量”。若要真支持“任意复向量”，需扩展相位加载逻辑。

    参数：
        x:        输入向量（目前：非负实）
        pos_init: 数据寄存器在整条线路中的起始物理位置（便于与更大电路拼接）

    返回：
        (circ, norm, tree, data_qubits)
            circ:        生成的量子电路
            norm:        原始向量的范数 ||x||
            tree:        幅度/相位二叉树
            data_qubits: 数据位索引列表（长度 n）
    """
    x = np.asarray(x)
    # === 重要限制（与文件头注释一致地说明）：当前实现仅接受“非负实向量” ===
    if not np.isrealobj(x):
        raise ValueError("Input vector must be real-valued.")
    if np.any(x < 0):
        raise ValueError("Input vector must be element-wise non-negative.")
    x = x.astype(float, copy=False)

    # 构建二叉树（内部会先把 x 归一化到 2^n 维）
    norm, B_x = create_tree_from_vector(x)
    n = int(np.ceil(np.log2(max(1, len(x)))))
    if n == 0:
        n = 1

    # 需要的辅助位数量：最多 (n-2)，不足则为 0
    ancilla_count = max(0, n - 2)
    total_qubits = n + ancilla_count
    circ = Circuit(total_qubits)

    # 数据位与辅助位的物理索引
    data_qubits = list(range(pos_init, pos_init + n))
    ancillas = list(range(pos_init + n, pos_init + n + ancilla_count))

    # 将数据位顺序镜像（把 LSB 放在物理右端），便于后续“深度=0 的位”就是 LSB
    for i in range(n // 2):
        circ.swap(pos_init + i, pos_init + n - 1 - i)

    # ===== 根节点旋转：在 LSB 上做一次 RY =====
    lsb = pos_init + (n - 1)
    B_l = B_x.get_left()
    B_r = B_x.get_right()
    B_l_is_leaf = B_l.get_left() is None
    node_value = B_l.get_value()[0] if B_l_is_leaf else B_l.get_value()
    theta0 = 2 * np.arccos(np.sqrt(node_value))
    circ.ry(lsb, float(theta0))

    # 根的相位处理：若是“复杂相位对”，补 RZ（相对相位）
    is_complex_pair = False
    if B_l_is_leaf:
        phase_l, phase_r = float(B_l.get_value()[1]), float(B_r.get_value()[1])
        if (phase_l, phase_r) not in [(0.0, 0.0), (0.0, np.pi), (np.pi, 0.0), (np.pi, np.pi)]:
            is_complex_pair = True
        else:
            # 离散 {0,π} 组合的情况已经被吸收到 theta0 的方向，根上无需再加 RZ
            pass
    if is_complex_pair:
        circ.rz(lsb, float(phase_r - phase_l))

    # ===== 递归处理左右子树 =====
    if n > 1:
        _process_node(circ, pos_init, n, B_l, ancillas)
        _process_node(circ, pos_init, n, B_r, ancillas)

    # 放一个整体 barrier，标识“数据寄存器已完成幅度装载”
    circ.barrier(*data_qubits)
    return circ, norm, B_x, data_qubits


# =============================================================================
# 五、（可选）仿真与评估工具：在本地有 cqlib 模拟器时用于验证
# =============================================================================
def simulate_statevector(circ: Circuit, as_array: bool = True):
    """
    试图用可用的 cqlib 模拟器拿到态矢量：
      优先 SimpleSimulator（纯 Python/Torch），退化用 StatevectorSimulator（C/C++）。
    返回：
      ndarray（复数 1D）或底层返回对象（当 as_array=False）
    """
    try:
        from cqlib.simulator import SimpleSimulator
        sim = SimpleSimulator(circ)
        sv = sim.statevector(dict_format=False)
        return np.array(sv, dtype=complex) if as_array else sv
    except Exception as e:
        try:
            from cqlib.simulator import StatevectorSimulator
            sim = StatevectorSimulator(circ, is_fusion=False)
            sv = sim.statevector()
            # 兼容某些返回 dict 的实现（键是二进制串）
            if isinstance(sv, dict):
                nq = int(np.log2(len(sv)))
                arr = np.zeros(2**nq, dtype=complex)
                for i in range(2**nq):
                    key = format(i, f'0{nq}b')
                    arr[i] = sv.get(key, 0.0)
                return arr if as_array else sv
            return np.array(sv, dtype=complex)
        except Exception as e2:
            raise RuntimeError(f"Could not simulate circuit: {e} | {e2}")


def normalize(v: np.ndarray) -> np.ndarray:
    """将向量标准化为单位范数（零向量时报错）。"""
    n = np.linalg.norm(v)
    if n == 0:
        raise ValueError("Zero vector not allowed")
    return v / n


def fidelity(psi: np.ndarray, target: np.ndarray) -> float:
    """F(|psi>, |target>) = |<target|psi>|^2（两者都会先做单位化）。"""
    psi = normalize(psi)
    tgt = normalize(target)
    return abs(np.vdot(tgt, psi))**2


def run_case(x: List[complex], name: str):
    """
    运行一个用例：
      1) 构造制备电路
      2) 仿真得到态矢量
      3) 构造“目标态”（数据寄存器态 ⊗ |0...0>_anc）
      4) 计算 Fidelity 并打印
    """
    print(f"\n=== Case: {name} | len={len(x)} ===")
    circ, norm, tree, data_qubits = prepare_state_from_vector(x)
    try:
        sv = simulate_statevector(circ, as_array=True)

        # 数据位数 n、辅助位数 anc
        n = int(np.ceil(np.log2(max(1, len(x))))) or 1
        anc = max(0, n - 2)

        # 目标：先把目标向量补到 2^n 维并归一化（仅数据寄存器）
        dim_data = 1 << n
        target_data = np.zeros(dim_data, dtype=complex)
        x_arr = np.asarray(x, dtype=complex)
        target_data[:len(x_arr)] = x_arr
        target_data /= np.linalg.norm(target_data)

        # 再嵌入到含 ancilla 的全空间（ancilla 在 MSB 且为 |0...0>）
        dim_full = 1 << (n + anc)
        target_full = np.zeros(dim_full, dtype=complex)
        target_full[:dim_data] = target_data

        F = fidelity(sv, target_full)
        print("模拟状态向量:", sv)
        print("目标状态:", target_full)
        print("Fidelity:", F)
        print("OK ✅" if F > 1 - 1e-10 else "Check ⚠️")
    except Exception as e:
        print("Simulation not available here. Circuit was generated. Error:", e)


# =============================================================================
# 六、简单自测
# =============================================================================
if __name__ == "__main__":
    # 注意：当前 prepare_state_from_vector 只接受“非负实向量”，
    # 下列用例若含负号/复相位，仅用于说明接口，运行时会触发 ValueError。
    run_case([1/np.sqrt(5), 2/np.sqrt(5)], "Basis |0>")
    run_case([0, 1], "Basis |1>")
    run_case([1, 0], "Real superposition (unnormalized)")
    run_case([1, 1], "Complex superposition (unnormalized)")  # 这里仍是实数；仅命名演示
    run_case([1, 0, 1, 0.5, 1, 0, 1, 0.5], "Length-4 real vector")
